from py3gpp import *
import numpy as np
import scipy.signal as signal

from NR.gscn import GSCN2CentFreq,gscn_pssIndices
from NR.pss import PSS
from Data.rawdata import RawData

class NR(RawData):
    def __init__(self,fn,mu=1,centered=False) -> None:
        RawData.__init__(self,fn)
        self.mu = mu
        self.centered = centered
        if not centered:
            self.dec_factor = int(self.fs/(512*15e3*2**self.mu))
        else:
            self.dec_factor = 1
        
        self.reset()
        
    def reset(self):
        self.carrier = nrCarrierConfig(NSizeGrid = 20, SubcarrierSpacing = 15 * 2**self.mu)
        self.info = nrOFDMInfo(self.carrier)
        self.Nfft = self.info['Nfft']

        slotGrid = nrResourceGrid(self.carrier)
        self.slotGrid = slotGrid[:, 0]
        self.dec_factor = int(self.fs/(self.Nfft*15e3 * 2**self.mu))
        
        self.detected_NID = -1
        
        self.max_frames = 4
    
    def set_search_end(self,up):
        self.f_end = up
        return self
    
    def set_search_start(self,down):
        self.f_start = down
        return self
    
    def read(self,start=5e-3,end=30e-3):
        RawData.read(self,start,end)
        if self.centered:
            self.waveform = self.raw
    
    def decimate(self,gscn):
        f_ssb = self.GSCN2CentFreq(gscn)
        RawData.decimate(self,f_ssb)
    
    def GSCN2CentFreq(self,gscn):
        self.gscn = gscn
        if gscn<7499:
            n = gscn % 3
            if n==0:
                N = gscn // 3
                M = 3
            elif n==1:
                N = gscn // 3
                M = 5
            else:
                N = gscn // 3 + 1
                M = 1
            return N*1.2e6+M*50e3
        elif gscn<22256:
            return (gscn-7499)*1.44e6+3e9
        else:
            return (gscn-22256)*17.28e6+24250.08e6
        
    @staticmethod
    def nearGSCN(fc):
        if fc<3e9:
            return int(round(fc/1.2e6*3))
        elif fc<24.250e9:
            return 7499+int(round((fc-3e9)/1.44e6))
        else:
            return 22256+int(round((fc-24250.08e6)/17.28))
        
    def report(self):
        r = f'{self.fn}:GSCN={self.gscn}  detect PSS {self.detected_NID2}@{self.timingOffset} with {self.peak}'
        if self.detected_NID!=-1:
            r += f' SSS {self.detected_NID}'
        return r
    
    def sync(self,SampleRate):
        peak_value = np.zeros(3)
        peak_index = np.zeros(3, 'int')
        PSS_LEN = 128
        NRE_PER_PRB = 12
        wv_cut = int(self.fs*25e-3)
        wv = self.waveform[:wv_cut]
        start = 56
        pssIndices = np.arange(start, start + PSS_LEN  - 1)
        # SampleRate = self.fs

        for current_NID2 in np.arange(3, dtype='int'):
            
            self.slotGrid[pssIndices] = nrPSS(current_NID2)
            [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)

            refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP

            temp = signal.correlate(wv, refWaveform, 'valid')  # correlate over 25 ms
            index = np.argmax(np.abs(temp))
            peak = np.abs(temp[index])
            peak_index[current_NID2] = index
            peak_value[current_NID2] = peak

        self.detected_NID2 = np.argmax(peak_value)
        self.timingOffset = peak_index[self.detected_NID2]
        self.peak =  peak_value[self.detected_NID2]
        
    def search(self,start=None,end=None):
        if start==None:
            start = self.f_start
        if end==None:
            end = self.f_end
        sg = NR.nearGSCN(start)
        se = NR.nearGSCN(end)
        n = se - sg
        offset = np.zeros((n,),dtype=np.int64)
        nid2 = np.zeros((n,),dtype=np.int32)
        peak = np.zeros((n,))
        for k in range(n):
            gscn = k+sg
            self.decimate(gscn)
            self.sync(self.fs/self.dec_factor)
            offset[k] = self.timingOffset
            peak[k] = self.peak
            nid2[k] = self.detected_NID2
        k = np.argmax(peak)
        self.gscn = k+sg
        self.detected_NID2 = nid2[k]
        self.peak = peak[k]
        self.timingOffset = offset[k]
        return k+sg,nid2[k],peak[k],offset[k]
    
    def OFDMModulate(
        carrier=None, grid=None, scs=None, initialNSlot=None, CyclicPrefix="normal", Nfft=None,
        SampleRate=None, Windowing=None, CarrierFrequency=0
    ):
        info = {}

        if carrier is None:
            if grid is None:
                print("Error: grid is needed when no carrierConfig is specified!")
                return
            NSizeGrid = grid.shape[0] // 12
            carrier = nrCarrierConfig(NSizeGrid=NSizeGrid)
            if initialNSlot is None:
                initialNSlot = carrier.NSlot
        else:
            NSizeGrid = carrier.NSizeGrid
            if initialNSlot is None:
                initialNSlot = carrier.NSlot
            
        if initialNSlot is None:
            initialNSlot = 0

        if scs is None:
            scs = carrier.SubcarrierSpacing

        if Nfft is None:
            if SampleRate is None:
                Nfft = nrOFDMInfo(nrb=NSizeGrid, scs=scs)["Nfft"]
                SampleRate = int(Nfft * scs * 1000)
            else:
                Nfft = int(SampleRate // scs // 1000)

        mu = (scs // 15) - 1

        info["Nfft"] = Nfft
        info["SampleRate"] = SampleRate
        info["CyclicPrefixLengths"] = np.empty(0, "int")
        info["SymbolLengths"] = np.empty(0, "int")
        if len(grid.shape) == 1:
            grid = grid[..., np.newaxis]
        nSlots = grid.shape[1]
        waveform = np.empty(0, "complex")
        if CyclicPrefix == "normal":
            N_cp1 = int(((144) * 2 ** (-mu) + 16) * (SampleRate / 30720000))
            N_cp2 = int((144 * 2 ** (-mu)) * (SampleRate / 30720000))
        else:
            N_cp1 = int((512 * 2 ** (-mu)) * (SampleRate / 30720000))
            N_cp2 = N_cp1
        N_cp = np.zeros(carrier.SymbolsPerSlot, dtype=int)
        for i in range(len(N_cp)):
            N_cp[i] = N_cp1 if i == 0 or i == 7 * 2 ** (mu) else N_cp2

        sample_pos_in_slot = 0
        for i in range(initialNSlot):
            sample_pos_in_slot += Nfft + N_cp[i]
            
        for sym_pos_in_grid in range(nSlots):
            sym_pos_in_slot = (sym_pos_in_grid + initialNSlot) % carrier.SymbolsPerSlot
            symbol_len = Nfft + N_cp[sym_pos_in_slot]
            nFill = (Nfft - grid.shape[0]) // 2
            if nFill < 0:
                full_slot_grid = grid[np.abs(nFill) :, sym_pos_in_grid][:nFill]
            else:
                full_slot_grid = np.concatenate([np.zeros(nFill), grid[:, sym_pos_in_grid], np.zeros(nFill)])

            # phase compensation according to TS 38.211 section 5.4
            if sym_pos_in_slot == 0:
                sample_pos_in_slot = 0
            sample_pos_in_slot += N_cp[sym_pos_in_slot]
            print(f'symbol {sym_pos_in_slot}, pos {sample_pos_in_slot}, CP {N_cp[sym_pos_in_slot]}, pos = {len(waveform)}')
            full_slot_grid *= np.exp(-1j * 2 * np.pi * CarrierFrequency / SampleRate * sample_pos_in_slot)
            sample_pos_in_slot += Nfft
            
            symbol_waveform = np.fft.ifft(np.fft.fftshift(full_slot_grid))
            symbol_waveform_cp = np.append(symbol_waveform[-N_cp[sym_pos_in_slot]:], symbol_waveform)
            waveform = np.append(waveform, symbol_waveform_cp)
            info["CyclicPrefixLengths"] = np.append(info["CyclicPrefixLengths"], N_cp[sym_pos_in_slot])
            info["SymbolLengths"] = np.append(info["SymbolLengths"], symbol_len)
            print(f'slot {sym_pos_in_slot}, cp_len {N_cp[sym_pos_in_slot]}')
        return [waveform, info]
